from matplotlib.figure import Figure
import numpy as np
import matplotlib
matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt

N=100
k1,k2=3,5
n = np.arange(N)

# 生成旋转因子矩阵
k_values = np.arange(max(k1,k2)+1)   # k = 0,1,2,3,4,5
k_column = k_values[:,None]  # 将k转为列向量
n_row = n[None,:]   # 将n转为行向量

# 离散傅里叶变换的旋转因子：W = e^(-j2πkn/N)
# 旋转因子矩阵 W 的维度：
# 行数：6行（k=0,1,2,3,4,5）
# 列数：100列（n=0,1,2,...,99）
# 每个元素 W[k,n] = cos(2πkn/N) - j·sin(2πkn/N) 表示：
# 在频率 k 下，时间点 n 的复指数信号
# 用于将时域信号分解为不同频率分量
W = np.exp(-2j * np.pi * k_column * n_row / N)

fig,axes = plt.subplots(1,3, figsize=(15,4))

# 实部：
# 当 k=0 时，实部恒为1（因为 e^0 = 1）
# 随着 k 增大，振荡频率增加

# 虚部：
# 当 k=0 时，虚部恒为0
# 虚部也呈现周期性变化

# 相位：
# 横轴 (n)：时间索引
# 纵轴 (k)：频率索引
# 颜色：表示复数的相位角（-π 到 π）
# 相位呈现线性变化
# 斜率与频率 k 成正比
titles = ["Re(W) 实部", "Im(W) 虚部", "Phase(W) 相位"]

data_list = [W.real, W.imag, np.angle(W)]

for ax, data, title in zip(axes, data_list, titles):
    im = ax.imshow(data, aspect = "auto", cmap = "twilight")
    ax.set_title(title)
    ax.set_xlabel("n (时间索引)")
    ax.set_ylabel("k (频率索引)")
    fig.colorbar(im, ax=ax)

plt.tight_layout()
plt.show()
z